# -*- coding: UTF-8 -*-

import preprocess
import utils
import gensim

class Similar_Word_Model():

    def __init__(self, use_article=False):
        self.save_file_name = 'sims_use_article.txt' if use_article else 'sims_not_use_article.txt'
        self.kv = {}
        self.load_sims()
        self.get_w2v_model(use_article)

    def save_sims(self):
        f = open(self.save_file_name, 'w', encoding='utf-8') 
        for k, v in self.kv.items():
            f.write(str(k)+' '+str(v)+'\n')
        f.close()

    def load_sims(self):
        f = open(self.save_file_name, 'r', encoding='utf-8')
        for line in f.readlines():
            line = line.strip()
            k = line.split(' ')[0]
            v = line.split(' ')[1]
            self.kv[k] = v
        f.close()

    def get_w2v_model(self, use_article=False):
        if use_article:
            article_json = utils.load_json('articles.json')
            sentences = []
            for article in article_json:
                if article['content']:
                    sentences.extend(preprocess.paragraph_preprocess(article['content']))
            self.model = gensim.models.Word2Vec(sentences, size=300)
            self.model.intersect_word2vec_format('GoogleNews-vectors-negative300.bin',
                                    lockf=1.0,
                                    binary=True)
            self.model.train(sentences, total_examples=model.corpus_count, epochs=model.epochs)
        else:
            self.model = gensim.models.KeyedVectors.load_word2vec_format('GoogleNews-vectors-negative300.bin', binary=True)
    
    def get_top_n_similar_word(self, word, topn=3):
        if word in self.kv:
            return self.kv[word]
        ret = []
        processed_word = preprocess.word_preprocess(word)
        # print('processed_word:',processed_word)
        try:
            arr = self.model.wv.most_similar(word)
        except KeyError:
            self.kv[word] = ret
            return ret
        else:
            for tu in arr:
                sim_word = tu[0]
                # print('sim_word:',sim_word)
                # if '_' in sim_word:
                #     continue
                # if re.search(r'[0-9]', sim_word):
                #     continue
                sim_value = tu[1]
                sim_word_preprocessed = preprocess.word_preprocess(sim_word)
                # print('sim_word_preprocessed:',sim_word_preprocessed)
                if sim_word_preprocessed != processed_word and sim_word_preprocessed not in ret:
                    ret.append(sim_word)
                if len(ret) >= topn:
                    break
            self.kv[word] = ret
            return ret